#!/usr/bin/env python
"""

    split_genome_files.py
    [--log_file PATH]
    [--quiet]

"""

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   options        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

from optparse import OptionParser
import sys, os
import StringIO
from commify import comma

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    sys.path.append(os.path.abspath(os.path.join(exe_path,"..", "python_modules")))
    module_name = os.path.split(sys.argv[0])[1]
    module_name = os.path.splitext(module_name)[0];
else:
    module_name = __name__

from options import check_mandatory_options


if __name__ == '__main__':
    parser = OptionParser(version="%prog 1.0", usage = "\n\n    %progs [options]")
    parser.add_option("-i", "--input_file", dest="input_file",
                      metavar="FILE", 
                      type="string",
                      help="name and path of genomic sequence data file in fasta format")
    parser.add_option("-o", "--output_path", dest="output_path",
                      metavar="PATTERN", 
                      type="string",
                      help="name and path of destination sequence files in fasta format." +
                            "e.g. '/data/h.{index:05d}.{contig}{contig_index}.fa' would produce " +
                            "'/data/h.00001.chrX.1.fa', '/data/h.00002.chrY.1.fa', ..."    )
    parser.add_option("-g", "--genomic_index", dest="genomic_index",
                        metavar="FILE", 
                        type="string",
                        help="name and path of index of genomic file")
    parser.add_option("-M", "--maximum_seq_len", dest="maximum_seq_len",
                        metavar="INT", 
                        type="int",
                        help="Number of bases in each segment")
    parser.add_option("-O", "--overlap", dest="overlap",
                        metavar="INT", 
                        type="int",
                        help="Number of bases successive segments overlap each other")
    
    
    
    #
    #   general options: verbosity / logging
    # 
    parser.add_option("-v", "--verbose", dest = "verbose",
                      action="store_true", default=False,
                      help="Print progress to STDERR")
    parser.add_option("-L", "--log_file", dest="log_file",
                      metavar="FILE", 
                      type="string",
                      help="Name and path of log file")
    parser.add_option("--skip_parameter_logging", dest="skip_parameter_logging",
                        action="store_true", default=False,
                        help="Do not print program parameters to log.")
    parser.add_option("-D", "--debug", dest="debug",
                        action="count", default=0,
                        help="Set default program parameters in debugging mode.")

    
    # get help string
    f =StringIO.StringIO()
    parser.print_help(f)
    helpstr = f.getvalue()
    (options, remaining_args) = parser.parse_args()
    
    if options.debug:
        test_dir = "/net/cpp-group/Leo/bin/src/python_modules/test/test_split_genomic_files/"
        if not options.output_path:
            options.output_path = test_dir + "test.{index:05d}.fa"
            #options.output_path = test_dir + "test.{index:05d}.{contig}.fa"
        if not options.input_file:
            options.input_file = test_dir + "oa_ornAna1_ensembl.fasta"
            options.genomic_index = test_dir + "oa_ornAna1_ensembl.idx"
            #options.input_file    = "/net/cpp-mirror/databases/indexed_fasta/hs_hg19_softmasked.fasta"
            #options.genomic_index = "/net/cpp-mirror/databases/indexed_fasta/hs_hg19_softmasked.idx"
        if not options.maximum_seq_len:
            #options.maximum_seq_len = 250000000
            options.maximum_seq_len = 200000000
        if not options.overlap:
            options.overlap = 100000
        options.log_file       = os.path.join(test_dir, "split_genome_files.log")
        options.verbose = 5
    
    
    # mandatory options
    mandatory_options = ["output_path", "input_file", "genomic_index", "maximum_seq_len", "overlap"]
    check_mandatory_options (options, mandatory_options, helpstr)
    if options.overlap > options.maximum_seq_len:
        raise Exception("--overlap should be smaller than maximum_seq_len" +
                     helpstr)


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   imports        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import StringIO
import re
import operator

from collections import defaultdict
from  get_genomic_sequence import get_genomic_sequence_coordinates_from_index_file, get_genomic_sequence
from sort_by_embedded_numbers import highlight_embedded_numbers
import logging
from lg_program_logging import setup_std_logging, MESSAGE
from options import get_option_strings

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

class output_file_state:
    def __init__ (self, output_pattern, max_seq_len):
        self.output_pattern     = output_pattern
        self.curr_file          = None
        self.curr_name          = None
        self.curr_contig        = None
        self.file_index         = 0
        self.per_contig_index   = 0
        self.max_seq_len        = max_seq_len
        self.prev_file_names    = set()
        self.len_seq_written    = 0

    def choose_output_file (self, contig, len_sequence_to_write, contig_pos, logger, cnt_contigs):
        """
        open new output file if current one too big
        """

        #
        # substitute contig name,
        #            per contig index
        #            per sequence file index
        if cnt_contigs < 30:
            logger.log(MESSAGE, "contig = %-20s, pos = %12s, len = %12s" % (contig, comma(contig_pos), comma(len_sequence_to_write)))
        if self.curr_contig != contig:
            self.per_contig_index = 0
            self.curr_contig = contig
            
        self.len_seq_written += len_sequence_to_write
        
        new_name = self.output_pattern.format(contig       = contig, 
                                              index        = self.file_index,
                                              per_contig_index = self.per_contig_index)
        # new file
        if self.curr_name != new_name or self.len_seq_written > self.max_seq_len:
            self.file_index += 1
            if self.curr_file:
                self.curr_file.close()

            new_name = self.output_pattern.format(contig       = contig, 
                                                  index        = self.file_index,
                                                  per_contig_index = self.per_contig_index)

            if new_name in self.prev_file_names:
                error_msg = (("Duplicate file name (%s). " % new_name) + 
                                "Try using {index} or "                + 
                                "{contig_index} in --output_path"      )
                logger.error(error_msg)
                raise Exception(error_msg)
            self.curr_name = new_name
            self.curr_file = open(self.curr_name, "w")
            file_name = os.path.split(self.curr_name)[1]
            logger.log(MESSAGE, "Opened new file %s for contig %-20s from %12s for %12s bp" % 
                            (file_name, self.curr_contig, comma(contig_pos), comma(len_sequence_to_write)))
            self.prev_file_names.add(self.curr_name)
            self.len_seq_written = len_sequence_to_write
            
        return self.curr_file


        
def output_split_sequence_file(sequence_file_name, index_file_name, output_pattern, max_seq_len, overlap_len, logger):
    sequence_loci = get_genomic_sequence_coordinates_from_index_file(open(index_file_name))

    cnt_contigs = len(sequence_loci)
    curr_state = output_file_state(output_pattern, max_seq_len)
#    import pprint
#    pp = pprint.PrettyPrinter(indent=4)
#    raise Exception("\n\n!!\n\n" + pp.pformat(sequence_loci) + "\n\n!!\n\n")
    for (contig, 
        (seq_offset, seq_len, ignore, acc_offset, acc_len)) in sorted(sequence_loci.iteritems(),
                                                key = lambda s:highlight_embedded_numbers(s[0])):

        estimate_seg_cnt = seq_len / (max_seq_len - overlap_len)
        if estimate_seg_cnt > 10000 :
            raise Exception ("Over 10000 files will be created just for contig %s. "
                             "Increase the value of --maximum_seq_len and reduce --overlap.")

        pos = 0
        per_contig_index = 0
        while pos < seq_len:
            segment_len = min(seq_len - pos, max_seq_len)
            
            curr_file = curr_state.choose_output_file (contig, segment_len, pos, logger, cnt_contigs)

            seq, new_contig = get_genomic_sequence (contig, pos, segment_len, 
                                                    sequence_loci, 
                                                    sequence_file_name)
            curr_file.write(">%s:%d:%d\n%s\n" % 
                                    (contig, pos, segment_len, seq))

            
            #print pos, max_seq_len, seq_len, segment_len, seq[0]+seq[-1]            
            if pos + max_seq_len >= seq_len:
                break
                        
            pos += max_seq_len - overlap_len

            
#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Main logic


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

if __name__ == '__main__':

    logger = logging.getLogger(module_name)
    setup_std_logging(logger, options.log_file, options.verbose)


    if not options.skip_parameter_logging:
        programme_name = os.path.split(sys.argv[0])[1]
        logger.info("%s %s" % (programme_name, " ".join(get_option_strings(parser, options))))

    if sys.argv.count("--debug"):
        sys.argv.remove("--debug")
        print options.genomic_index, options.input_file, options.output_path, options.maximum_seq_len, options.overlap
        # unittest.main()
        output_split_sequence_file(options.input_file, options.genomic_index, options.output_path, options.maximum_seq_len, options.overlap, logger)
    else:
        output_split_sequence_file(options.input_file, options.genomic_index, options.output_path, options.maximum_seq_len, options.overlap, logger)
    
    
